import csv
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Patch  # 新增：用于创建图例代理对象
from mpl_toolkits.mplot3d import Axes3D
from scipy.stats import norm

# 全局计数器：控制表头打印次数（仅前5次）
header_print_count = 0

class SimulationConfig:
    """模拟配置类，管理所有物理参数"""
    def __init__(self, config_dict=None):
        if config_dict is None:
            # 默认配置
            self.W_mix = 20  # 混合气体平均电离能 (eV/电荷)
            self.Fano_factor = 0.3  # Fano因子
            self.GMCP_gain = 10**3  # 电子倍增因子
            self.hole_efficiency = 0.96  # 孔道入口效率
            self.collection_efficiency = 0.11  # TopMetal-II的收集率
            self.v_drift = 0.34  # 纵向漂移速度 (mm/ns) @ 1kV/cm电场强度
            self.bottom_adsorption = 0.26  # 底部吸附率
            self.MCP_Z_POSITION = -6.9  # MCP位于Z=-6.9mm平面
            self.TOPMETAL_Z_POSITION = self.MCP_Z_POSITION - 3.4  # TopMetal位于MCP下方3.4mm处
            self.DRIFT_REGION_DIFFUSION = 0.02066865  # 漂移区横向扩散系数 (mm/√mm)
            self.LONGITUDINAL_DIFFUSION = 0.02066865  # 漂移区纵向扩散系数 (mm/√mm)
            self.INDUCTION_REGION_E_FIELD_RATIO = 2.0  # 感应区电场强度是漂移区的2倍
            self.SAMPLING_INTERVAL = 5  # 采样间隔 (ns)
            self.VERTICAL_DIFFUSION = 0.1  # 纵向均匀扩散系数 (mm)
            # 体元离散化参数
            self.VOXEL_SIZE = 0.045  # 体元大小，单位mm (45μm = 0.045mm)
            self.VOXEL_DIMENSIONS = (72, 72, 72)  # 体元网格尺寸
            # 新增：三维体素尺寸（x, y, z方向统一）
            self.voxel_size = (self.VOXEL_SIZE, self.VOXEL_SIZE, self.VOXEL_SIZE)
        else:
            # 从字典加载配置
            self.__dict__.update(config_dict)
            # 确保voxel_size存在（若配置字典未包含，设为默认值）
            if not hasattr(self, 'voxel_size'):
                self.voxel_size = (self.VOXEL_SIZE, self.VOXEL_SIZE, self.VOXEL_SIZE) if hasattr(self, 'VOXEL_SIZE') else (0.045, 0.045, 0.045)
    
    def get_diffusion_model(self):
        """创建并返回扩散模型实例"""
        return DiffusionModel(self)

class DiffusionModel:
    """抽象扩散模型基类"""
    def __init__(self, config):
        self.config = config
    
    def calculate_drift_distance(self, position_mm, target_z_mm):
        """计算电子到目标平面的漂移距离(单位:mm)"""
        z_position_mm = position_mm[2]
        drift_distance_mm = abs(z_position_mm - target_z_mm)
        return drift_distance_mm 
    
    def calculate_drift_sigma(self, position):
        """计算漂移过程的扩散标准差"""
        drift_distance_mm = self.calculate_drift_distance(position, self.config.MCP_Z_POSITION)
        transverse_sigma_mm = self.config.DRIFT_REGION_DIFFUSION * np.sqrt(drift_distance_mm)
        return transverse_sigma_mm 
    
    def calculate_induction_sigma(self):
        """计算感应过程的扩散标准差"""
        drift_distance_mm = self.calculate_drift_distance(
            (0, 0, self.config.MCP_Z_POSITION), 
            self.config.TOPMETAL_Z_POSITION
        )
        # 扩散系数与电场强度的平方根成反比
        adjusted_diffusion = self.config.DRIFT_REGION_DIFFUSION / np.sqrt(self.config.INDUCTION_REGION_E_FIELD_RATIO)
        transverse_sigma_mm = adjusted_diffusion * np.sqrt(drift_distance_mm)
        return transverse_sigma_mm
    
    def get_induction_drift_velocity(self):
        """获取感应区的漂移速度（电场强度加倍，速度加倍）"""
        return self.config.v_drift * self.config.INDUCTION_REGION_E_FIELD_RATIO
    
    def calculate_longitudinal_drift_sigma(self, position):
        """计算漂移过程的纵向扩散标准差"""
        drift_distance_mm = self.calculate_drift_distance(position, self.config.MCP_Z_POSITION)
        longitudinal_sigma_mm = self.config.LONGITUDINAL_DIFFUSION * np.sqrt(drift_distance_mm)
        return longitudinal_sigma_mm
    
    def calculate_longitudinal_induction_sigma(self):
        """计算感应过程的纵向扩散标准差"""
        drift_distance_mm = self.calculate_drift_distance(
            (0, 0, self.config.MCP_Z_POSITION), 
            self.config.TOPMETAL_Z_POSITION
        )
        # 扩散系数与电场强度的平方根成反比
        adjusted_diffusion = self.config.LONGITUDINAL_DIFFUSION / np.sqrt(self.config.INDUCTION_REGION_E_FIELD_RATIO)
        longitudinal_sigma_mm = adjusted_diffusion * np.sqrt(drift_distance_mm)
        return longitudinal_sigma_mm
    
    def apply_diffusion(self, edep_position):
        """应用扩散模型，返回最终位置（考虑纵向扩散）"""
        drift_sigma = self.calculate_drift_sigma(edep_position)
        induction_sigma = self.calculate_induction_sigma()
        # 使用单独的纵向扩散系数计算纵向扩散标准差
        longitudinal_drift_sigma = self.calculate_longitudinal_drift_sigma(edep_position)
        longitudinal_induction_sigma = self.calculate_longitudinal_induction_sigma()
        
        # 第一阶段扩散（沉积→MCP）
        mcp_x = edep_position[0] + np.random.normal(0, drift_sigma)
        mcp_y = edep_position[1] + np.random.normal(0, drift_sigma)
        mcp_z = self.config.MCP_Z_POSITION + np.random.normal(0, longitudinal_drift_sigma)
        
        # 第二阶段扩散（MCP→Topmetal）
        topmetal_x = mcp_x + np.random.normal(0, induction_sigma)
        topmetal_y = mcp_y + np.random.normal(0, induction_sigma)
        topmetal_z = self.config.TOPMETAL_Z_POSITION + np.random.normal(0, longitudinal_induction_sigma)
        
        # 计算横向偏移量
        delta_x_mcp = mcp_x - edep_position[0]
        delta_y_mcp = mcp_y - edep_position[1]
        delta_x_topmetal = topmetal_x - mcp_x
        delta_y_topmetal = topmetal_y - mcp_y
        
        # 最终位置（考虑纵向扩散）
        final_x = edep_position[0] + delta_x_mcp + delta_x_topmetal
        final_y = edep_position[1] + delta_y_mcp + delta_y_topmetal
        # 直接使用经过纵向扩散的z位置
        final_z = topmetal_z
        
        # 纵向扩散总量（用于统计分析）
        total_longitudinal_diffusion = (mcp_z - self.config.MCP_Z_POSITION) + (topmetal_z - self.config.TOPMETAL_Z_POSITION)
        
        return final_x, final_y, final_z, total_longitudinal_diffusion, self.calculate_arrival_time(edep_position)
    
    def calculate_arrival_time(self, edep_position):
        """计算电荷到达时间"""
        t_drift_to_mcp = abs(edep_position[2] - self.config.MCP_Z_POSITION) / self.config.v_drift
        t_drift_to_topmetal = abs(self.config.MCP_Z_POSITION - self.config.TOPMETAL_Z_POSITION) / self.get_induction_drift_velocity()
        return t_drift_to_mcp + t_drift_to_topmetal

def read_deposits_info(csv_path):
    """读取CSV文件中的能量沉积事件（仅打印前5个表头）"""
    global header_print_count  # 引用全局计数器
    deposits = []
    with open(csv_path, 'r') as f:
        reader = csv.reader(f)
        in_deposits = False
        header = None
        for row in reader:
            row = [cell.strip() for cell in row]
            if not row:
                continue
            if "Deposits Info" in row:
                in_deposits = True
                try:
                    header = next(reader)
                    # 仅打印前5个表头信息
                    if header_print_count < 5:
                        print(f"找到Deposits Info,表头字段：{header}")
                        header_print_count += 1  # 计数器递增
                except StopIteration:
                    print("错误：存款信息部分缺少表头")
                    return []
                continue
            if in_deposits and header:
                try:
                    pre_x = float(row[header.index('pre_pos_x')])
                    pre_y = float(row[header.index('pre_pos_y')])
                    pre_z = float(row[header.index('pre_pos_z')])
                    post_x = float(row[header.index('post_pos_x')])
                    post_y = float(row[header.index('post_pos_y')])
                    post_z = float(row[header.index('post_pos_z')])
                    edep_MeV = float(row[header.index('edep')])
                    
                    deposits.append({
                        'pre_pos_x': pre_x,
                        'pre_pos_y': pre_y,
                        'pre_pos_z': pre_z,
                        'post_pos_x': post_x,
                        'post_pos_y': post_y,
                        'post_pos_z': post_z,
                        'edep_MeV': edep_MeV
                    })
                except ValueError as e:
                    print(f"数据解析错误：{e}，跳过该行数据：{row}")
                except IndexError as e:
                    print(f"表头字段缺失：{e}，跳过该行数据：{row}")
    return deposits

def process_deposits(deposits, config=None):
    """处理能量沉积事件"""
    if config is None:
        config = SimulationConfig()
    
    diffusion_model = config.get_diffusion_model()
    charge_positions = []
    
    for dep in deposits:
        if dep['edep_MeV'] <= 0:
            continue
            
        mid_x = (dep['pre_pos_x'] + dep['post_pos_x']) / 2
        mid_y = (dep['pre_pos_y'] + dep['post_pos_y']) / 2
        mid_z = (dep['pre_pos_z'] + dep['post_pos_z']) / 2
        edep_position = (mid_x, mid_y, mid_z)
        
        # 计算原始电离电荷数
        edep_eV = dep['edep_MeV'] * 1e6
        mean_ion = edep_eV / config.W_mix
        sigma_ion = np.sqrt(config.Fano_factor * mean_ion)
        n_ion = max(1, int(np.round(norm.rvs(loc=mean_ion, scale=sigma_ion))))
        
        # 应用GMCP放大效应
        effective_ion = n_ion * config.hole_efficiency * (1 - config.bottom_adsorption)
        amplified_charge = int(effective_ion * config.GMCP_gain * config.collection_efficiency)
        
        # 为每个电荷应用扩散（考虑纵向扩散）
        for _ in range(amplified_charge):
            final_x, final_y, final_z, long_diff, t_total = diffusion_model.apply_diffusion(edep_position)
            charge_positions.append( (final_x, final_y, final_z, long_diff, t_total, edep_position[2]) )
    
    return np.array(charge_positions) if charge_positions else np.array([])

def sample_trajectory(charge_positions, config=None):
    """对径迹进行采样并添加纵向扩散"""
    if config is None:
        config = SimulationConfig()
    
    if len(charge_positions) == 0:
        print("警告:无电荷数据用于采样")
        return np.array([])
    
    # 提取数据
    x, y, z, long_diff, arrival_time, original_z = charge_positions.T
    
    # 找到径迹的最高点和最低点
    max_z = np.max(z)
    min_z = np.min(z)
    
    # 计算对应的时间范围（假设Z方向速度恒定）
    max_time = arrival_time[np.argmax(z)]
    min_time = arrival_time[np.argmin(z)]
    
    # 生成采样时间点（从径迹最上面开始，每隔5ns采样一次）
    sample_times = np.arange(min_time, max_time + config.SAMPLING_INTERVAL, config.SAMPLING_INTERVAL)
    
    # 为每个采样时间点创建一个采样集
    sampled_positions = []
    
    for t in sample_times:
        # 选择在该时间点附近的电荷
        mask = np.abs(arrival_time - t) < config.SAMPLING_INTERVAL / 2
        if not np.any(mask):
            continue
            
        # 提取这些电荷的位置
        sample_x = x[mask]
        sample_y = y[mask]
        sample_original_z = original_z[mask]  # 使用原始沉积位置的Z坐标
        
        # 为每个电荷添加纵向均匀扩散 - 不再与时间相关
        diffusion_amount = config.VERTICAL_DIFFUSION
        diffused_z = sample_original_z + np.random.normal(0, diffusion_amount, size=len(sample_x))
        
        # 添加到采样后的径迹中
        for i in range(len(sample_x)):
            sampled_positions.append( (sample_x[i], sample_y[i], diffused_z[i], t) )
    
    return np.array(sampled_positions) if sampled_positions else np.array([])

def plot_comparison(original_positions, sampled_positions, config=None):
    """绘制采样前后的三维电荷分布对比图（增强纵向扩散可视化）"""
    if config is None:
        config = SimulationConfig()
    
    if len(original_positions) == 0 or len(sampled_positions) == 0:
        print("警告:无有效电荷数据用于绘图")
        return
    
    # 提取原始电荷的三维坐标 (x, y, z)
    orig_x, orig_y, orig_z = original_positions[:, 0], original_positions[:, 1], original_positions[:, 2]
    
    # 提取采样后电荷的三维坐标 (x, y, z)
    samp_x, samp_y, samp_z = sampled_positions[:, 0], sampled_positions[:, 1], sampled_positions[:, 2]
    
    # 创建对比图（左右布局）
    fig = plt.figure(figsize=(24, 10))
    
    # -------------------------- 左侧：采样前的三维电荷分布 --------------------------
    ax1 = fig.add_subplot(121, projection='3d')
    
    # 关键优化1：用颜色映射表示Z轴位置（直观区分纵向分布）
    sc1 = ax1.scatter(
        orig_x, orig_y, orig_z,
        s=3,                  # 增大点大小，突出Z轴分布
        alpha=0.6,            # 适度透明度，避免完全遮挡
        c=orig_z,             # 颜色随Z轴坐标变化
        cmap='coolwarm',      # 冷暖色区分Z轴高低
        edgecolors='none',
        label='原始电荷'
    )
    # 添加Z轴颜色条
    cbar1 = plt.colorbar(sc1, ax=ax1, shrink=0.8)
    cbar1.set_label('Z坐标 (mm)', fontsize=9)
    
    # 绘制参考平面
    x_range = np.linspace(min(orig_x.min(), samp_x.min()) - 1, max(orig_x.max(), samp_x.max()) + 1, 20)
    y_range = np.linspace(min(orig_y.min(), samp_y.min()) - 1, max(orig_y.max(), samp_y.max()) + 1, 20)
    xx, yy = np.meshgrid(x_range, y_range)
    ax1.plot_surface(
        xx, yy, np.full_like(xx, config.MCP_Z_POSITION),
        color='lightcoral', alpha=0.3, zorder=5
    )
    ax1.plot_surface(
        xx, yy, np.full_like(xx, config.TOPMETAL_Z_POSITION),
        color='cyan', alpha=0.3, zorder=5
    )
    
    # 关键优化2：坐标轴等比例缩放（X/Y/Z单位长度视觉一致）
    ax1.set_box_aspect([1, 1, 1])  # 三个轴的显示比例相同
    # 关键优化3：调整视角（侧视+倾斜，突出Z轴分布）
    ax1.view_init(elev=20, azim=120)
    
    # 图例与标签
    from matplotlib.patches import Patch
    mcp_proxy = Patch(color='lightcoral', alpha=0.3, label='MCP平面')
    topmetal_proxy = Patch(color='cyan', alpha=0.3, label='Topmetal平面')
    ax1.legend(handles=[sc1, mcp_proxy, topmetal_proxy], loc='upper right', fontsize=9)
    ax1.set_xlabel('X (mm)', fontsize=10)
    ax1.set_ylabel('Y (mm)', fontsize=10)
    ax1.set_zlabel('Z (mm)', fontsize=10)
    ax1.set_title('采样前的三维电荷分布（颜色=Z坐标）', fontsize=12, pad=20)
    
    # -------------------------- 右侧：采样后的三维电荷分布 --------------------------
    ax2 = fig.add_subplot(122, projection='3d')
    
    # 同样用颜色映射表示Z轴位置
    sc2 = ax2.scatter(
        samp_x, samp_y, samp_z,
        s=3,
        alpha=0.6,
        c=samp_z,
        cmap='coolwarm',
        edgecolors='none',
        label='采样后电荷'
    )
    cbar2 = plt.colorbar(sc2, ax=ax2, shrink=0.8)
    cbar2.set_label('Z坐标 (mm)', fontsize=9)
    
    # 绘制参考平面
    ax2.plot_surface(
        xx, yy, np.full_like(xx, config.MCP_Z_POSITION),
        color='lightcoral', alpha=0.3, zorder=5
    )
    ax2.plot_surface(
        xx, yy, np.full_like(xx, config.TOPMETAL_Z_POSITION),
        color='cyan', alpha=0.3, zorder=5
    )
    
    # 坐标轴等比例+优化视角
    ax2.set_box_aspect([1, 1, 1])
    ax2.view_init(elev=20, azim=120)
    
    # 图例与标签
    ax2.legend(handles=[sc2, mcp_proxy, topmetal_proxy], loc='upper right', fontsize=9)
    ax2.set_xlabel('X (mm)', fontsize=10)
    ax2.set_ylabel('Y (mm)', fontsize=10)
    ax2.set_zlabel('Z (mm)', fontsize=10)
    ax2.set_title('采样后的三维电荷分布（颜色=Z坐标）', fontsize=12, pad=20)
    
    # 调整布局
    plt.tight_layout()
    plt.subplots_adjust(top=0.9)
    plt.show()
    
def validate_diffusion_results(positions, deposits, config=None):
    """验证扩散结果的统计特性（考虑纵向扩散）"""
    if config is None:
        config = SimulationConfig()
    
    if len(positions) == 0:
        print("无有效电荷数据用于验证")
        return
    
    # 计算所有沉积点的平均位置
    edep_centers = np.array([
        [(d['pre_pos_x'] + d['post_pos_x'])/2, 
         (d['pre_pos_y'] + d['post_pos_y'])/2, 
         (d['pre_pos_z'] + d['post_pos_z'])/2] 
        for d in deposits
    ])
    
    # 计算扩散后的电荷分布统计
    x, y, z, long_diff, arrival_time, _ = positions.T
    charge_mean = np.mean([x, y, z], axis=1)
    charge_std = np.std([x, y, z], axis=1)
    
    # 计算预期的横向扩散标准差
    avg_drift_distance = np.mean([
        calculate_drift_distance(center, config.MCP_Z_POSITION)
        for center in edep_centers
    ])
    
    induction_distance = calculate_drift_distance(
        (0, 0, config.MCP_Z_POSITION), 
        config.TOPMETAL_Z_POSITION
    )
    
    # 考虑电场差异的扩散系数
    drift_diffusion = config.DRIFT_REGION_DIFFUSION
    induction_diffusion = drift_diffusion / np.sqrt(config.INDUCTION_REGION_E_FIELD_RATIO)
    
    # 横向扩散标准差
    expected_transverse_sigma = np.sqrt(
        drift_diffusion**2 * avg_drift_distance +
        induction_diffusion**2 * induction_distance
    )
    
    # 使用单独的纵向扩散系数计算预期的纵向扩散标准差
    longitudinal_drift_diffusion = config.LONGITUDINAL_DIFFUSION
    longitudinal_induction_diffusion = longitudinal_drift_diffusion / np.sqrt(config.INDUCTION_REGION_E_FIELD_RATIO)
    
    expected_longitudinal_sigma = np.sqrt(
        longitudinal_drift_diffusion**2 * avg_drift_distance +
        longitudinal_induction_diffusion**2 * induction_distance
    )
    
    # 计算实际纵向扩散统计
    long_diff_mean = np.mean(long_diff)
    long_diff_std = np.std(long_diff)
    
    print(f"\n电荷分布统计:")
    print(f"  平均位置: ({charge_mean[0]:.3f}, {charge_mean[1]:.3f}, {charge_mean[2]:.3f}) mm")
    print(f"  横向标准差: (X: {charge_std[0]:.3f}, Y: {charge_std[1]:.3f}) mm")
    print(f"  预期横向标准差: {expected_transverse_sigma:.3f} mm")
    print(f"  纵向扩散统计:")
    print(f"    平均纵向扩散: {long_diff_mean:.3f} mm")
    print(f"    纵向扩散标准差: {long_diff_std:.3f} mm")
    print(f"    预期纵向标准差: {expected_longitudinal_sigma:.3f} mm")
    
    # 打印电场和速度信息
    induction_velocity = config.v_drift * config.INDUCTION_REGION_E_FIELD_RATIO
    print(f"\n电场和速度信息:")
    print(f"  漂移区电场强度: 1.0 (相对单位)")
    print(f"  感应区电场强度: {config.INDUCTION_REGION_E_FIELD_RATIO} (相对单位)")
    print(f"  漂移区漂移速度: {config.v_drift:.3f} mm/ns")
    print(f"  感应区漂移速度: {induction_velocity:.3f} mm/ns")

def calculate_drift_distance(position, target_z):
    """辅助函数：计算漂移距离（单位：mm）"""
    z_position = position[2]
    drift_distance_mm = abs(z_position - target_z)
    return drift_distance_mm

def voxelize_electrons(electron_positions, config=None):
    """将电子位置离散化为体元网格"""
    if config is None:
        config = SimulationConfig()
        
    if len(electron_positions) == 0:
        print("警告: 无电子数据用于体元离散化")
        return np.zeros(config.VOXEL_DIMENSIONS), (0, 0, 0)
        
    # 提取电子位置（确保只取前三维x,y,z）
    if len(electron_positions.shape) == 1:
        x = np.array([electron_positions[0]])
        y = np.array([electron_positions[1]])
        z = np.array([electron_positions[2]])
    else:
        x, y, z = electron_positions[:, 0], electron_positions[:, 1], electron_positions[:, 2]
        
    # 计算电子云的边界和中心
    x_min, x_max = np.min(x), np.max(x)
    y_min, y_max = np.min(y), np.max(y)
    z_min, z_max = np.min(z), np.max(z)
    
    center_x = (x_min + x_max) / 2
    center_y = (y_min + y_max) / 2
    center_z = (z_min + z_max) / 2
    
    # 计算网格的尺寸（以mm为单位）
    grid_size_mm = config.VOXEL_SIZE * config.VOXEL_DIMENSIONS[0]
    
    # 设置网格原点，使电子云中心位于网格中心
    offset_x = center_x - grid_size_mm / 2
    offset_y = center_y - grid_size_mm / 2
    offset_z = center_z - grid_size_mm / 2
    offset = (offset_x, offset_y, offset_z)
    
    # 创建体元网格
    voxel_grid = np.zeros(config.VOXEL_DIMENSIONS, dtype=int)
    
    # 将每个电子分配到对应的体元
    for i in range(len(x)):
        # 计算电子在网格中的索引
        voxel_x = int((x[i] - offset_x) / config.VOXEL_SIZE)
        voxel_y = int((y[i] - offset_y) / config.VOXEL_SIZE)
        voxel_z = int((z[i] - offset_z) / config.VOXEL_SIZE)
        
        # 检查索引是否在网格范围内
        if (0 <= voxel_x < config.VOXEL_DIMENSIONS[0] and 
            0 <= voxel_y < config.VOXEL_DIMENSIONS[1] and 
            0 <= voxel_z < config.VOXEL_DIMENSIONS[2]):
            voxel_grid[voxel_x, voxel_y, voxel_z] += 1
            
    return voxel_grid, offset

def plot_voxelized_data(voxel_grid, offset, config=None):
    """可视化体元化后的电子分布"""
    if config is None:
        config = SimulationConfig()
        
    fig = plt.figure(figsize=(12, 10))
    ax = fig.add_subplot(111, projection='3d')
    
    # 获取有电子的体元坐标
    x_coords, y_coords, z_coords = np.where(voxel_grid > 0)
    counts = voxel_grid[x_coords, y_coords, z_coords]
    
    # 将体元索引转换回实际坐标（mm）
    actual_x = x_coords * config.VOXEL_SIZE + offset[0]
    actual_y = y_coords * config.VOXEL_SIZE + offset[1]
    actual_z = z_coords * config.VOXEL_SIZE + offset[2]
    
    # 绘制体元化后的电子分布
    scatter = ax.scatter(actual_x, actual_y, actual_z, 
                         c=counts, cmap='RdYlBu_r', 
                         s=counts*3, alpha=0.7)
    
    # 添加颜色条和标签
    cbar = plt.colorbar(scatter, ax=ax)
    cbar.set_label('体元内电子数')
    
    ax.set_xlabel('X (mm)')
    ax.set_ylabel('Y (mm)')
    ax.set_zlabel('Z (mm)')
    ax.set_title(f'体元化后的电子分布 ({config.VOXEL_DIMENSIONS[0]}×{config.VOXEL_DIMENSIONS[1]}×{config.VOXEL_DIMENSIONS[2]}, 体元大小: {config.VOXEL_SIZE*1000}μm)')
    
    plt.tight_layout()
    plt.show()

# 主程序入口
if __name__ == "__main__":
    csv_path = "/home/qht/Re_Track/Track_data//YDM/event_0.csv"  # 请确认路径正确性
    print(f"正在读取数据文件：{csv_path}")
    
    # 读取能量沉积事件
    deposits = read_deposits_info(csv_path)
    if not deposits:
        print("错误：未读取到任何能量沉积信息")
        exit(1)
    
    # 配置模拟参数
    config = SimulationConfig()
    
    # 处理数据并应用扩散
    print("\n处理能量沉积事件...")
    charge_positions = process_deposits(deposits, config)
    print(f"生成电荷总数: {len(charge_positions)}")
    
    # 验证扩散结果
    validate_diffusion_results(charge_positions, deposits, config)
    
    # 对径迹进行采样并添加纵向扩散
    print("\n对径迹进行采样并添加纵向扩散...")
    sampled_positions = sample_trajectory(charge_positions, config)
    print(f"采样后电荷总数: {len(sampled_positions)}")
    
    # 将采样后的电子进行体元离散化
    print("\n将电子进行体元离散化...")
    voxel_grid, offset = voxelize_electrons(sampled_positions[:, :3], config)  # 只使用前三个维度(x,y,z)
    active_voxels = np.sum(voxel_grid > 0)
    print(f"非零体元数量: {active_voxels}")
    print(f"体元网格大小: {config.VOXEL_DIMENSIONS[0]}×{config.VOXEL_DIMENSIONS[1]}×{config.VOXEL_DIMENSIONS[2]}")
    print(f"体元大小: {config.VOXEL_SIZE*1000}μm")
    
    # 可视化体元化后的数据
    plot_voxelized_data(voxel_grid, offset, config)
    
    # 绘制采样前后的径迹对比图
    print("\n绘制采样前后的径迹对比图...")
    plot_comparison(charge_positions, sampled_positions, config)